import numpy as np
import time
import functools
import os
import re
import gc
import logging

import labrad
from calibrator.ghzdac import calibrate
from calibrator.ghzdac import keys
from calibrator.servers.anritsu_MG3692C import AnritsuServer
from calibrator.servers.zurich_fpga import ziDAQ


####################################################################
# dataprocess tools                                                #
####################################################################

exp_data_dir = r'M:\Experimental Data'
calibration_session = 'Zurich Calibration'


def get_data_path(board_name: str, awg_index: int):
    _dir_list = [calibration_session, board_name, 'awg' + str(awg_index)]

    _dir_list_dir = list(map(lambda x: x+str(r'.dir'), _dir_list))
    return os.path.join(exp_data_dir, *_dir_list_dir)


def _get_dataset(idx: int, path: str):
    """
    return the data from the file name matching the parttern
    Example: "path\00010 - example.csv" (idx = 10)
    Returns:
        data (numpy.ndarray)
    """
    idx_length = len(str(idx))
    if idx_length > 5:
        raise ValueError("idx is too large, should < 99999", idx)

    re_parttern = r"^0{%d}%d - .*\.csv" % (5-idx_length, idx)
    def re_func(x): return re.match(re_parttern, x)
    names = list(fn for fn in os.listdir(path) if re_func(fn))

    if len(names) is not 1:
        raise Exception("regex match wrong file paths", names)

    dataset_name = os.path.join(path, names[0])
    data = np.loadtxt(dataset_name, delimiter=',')
    print("loaded", dataset_name)
    return data


def get_dataset(idx: int, board_name: str, awg_index: int):
    """
    return the data from specific board_name and awg_index
    """
    path = get_data_path(board_name, awg_index)
    data = _get_dataset(idx, path)
    return data

def get_now_time():
    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    
####################################################################
# labrad tools                                                     #
####################################################################

def add_reg_parameters(reg, dv):
    para_keys = reg._get_list()[1]
    for para_key in para_keys:
        _paras = (para_key, str(reg[para_key] ))
        dv.add_parameter(_paras)
    return




####################################################################
# DAC zero calibration                                             #
####################################################################

def _cali_zero(cxn, LO_Server, dv, reg, fpga,
    freq=np.arange(4.0, 7.001, 0.005), noisy=True, comment='',
    max_reps = 3, limit_power = -80.):
    name = 'zero'
    dv.new(name, [('Frequency', 'GHz')], ['I', 'Q', 'power cal'])
    add_reg_parameters(reg, dv)
    _now_time = str(get_now_time())
    paras = [('Time', _now_time), ('comment', comment)]
    for para in paras:
        dv.add_parameter(para)
    for _freq in freq:
        now_reps = 0
        while now_reps < max_reps:
            now_reps += 1
            daczeros = calibrate.zeroFixedCarrier(cxn, LO_Server, reg, fpga, _freq)
            power = daczeros[-1]
            if power < limit_power:
                now_reps = max_reps
                
        dv_send = np.hstack([_freq, daczeros])
        if noisy:
            print(dv_send)
        dv.add(dv_send)
    return


def save_zero(dv, reg, data, comment=''):
    name = 'zero'
    dv.new(name, [('Frequency', 'GHz')], ['I', 'Q', 'power cal'])
    add_reg_parameters(reg, dv)
    _now_time = str(get_now_time())
    paras = [('Time', _now_time), ('comment', comment)]
    for para in paras:
        dv.add_parameter(para)
    dv.add(data)
    return


class user_interface(object):
    def __init__(self):
        self.cxn = labrad.connect()
        try:
            dv = self.cxn.data_vault
            self.dv = dv
        except Exception:
            raise("labrad server data_vault not exist")

        self.fpga = ziDAQ()
        self.board_name = 'dev8334'
        self.awg_index = 2
        self.fpga.set_port(self.board_name, self.awg_index)
        self.dv.cd(['', calibration_session, self.board_name,
            'awg'+str(self.awg_index)])

        self.reg = keys.new_registry(self.cxn,
            board_name= self.board_name, use_default=False)

        self.LO_Server = None
        self.servers = (self.cxn, self.LO_Server, self.dv,
                        self.reg, self.fpga)

    def _set_default_reg(self, board_name: str):
        """
        set default parameters in registry.
        only run once for a given device
        """
        keys.new_registry(self.cxn, board_name=board_name, use_default=True)

    def _make_datavault_subdir(self, board_name: str, awg_indexs: list or None = None):
        """create new folders for data vault
        only run once for a given device
        """
        self.dv.cd(['', 'Zurich Calibration'])
        try:
            self.dv.mkdir(board_name)
        except Exception: pass

        if not awg_indexs:
            awg_indexs = [0, 1, 2, 3]

        self.dv.cd(['', 'Zurich Calibration', board_name])
        for awg_index in awg_indexs:
            subdir_name = 'awg'+str(awg_index)
            self.dv.mkdir(subdir_name)
        return

    def add_new_device(self, board_name: str):
        logging.warning("These run only at first time for a given device")
        self._set_default_reg(board_name)
        self._make_datavault_subdir(board_name, awg_indexs=[0,1,2,3])

    def update_LO(self, LO_id: str = None):
        if LO_id is None:
            LO_ids = [self.reg['LO_id']]
        else:
            LO_ids = [LO_id]
        
        # clear visa Memory
        del self.LO_Server
        gc.collect()
        
        self.LO_Server = AnritsuServer(LO_ids)


    def switch_port(self, board_name, awg_index):
        """switch to a new port of zurich devices
        """
        self.board_name = board_name
        self.awg_index = awg_index
        self.fpga.set_port(board_name, awg_index)
        self.dv.cd(['', 'Zurich Calibration', board_name, 'awg'+str(awg_index)])
        self.reg = keys.new_registry(self.cxn,
            board_name= board_name, use_default=False)

    def cali_zero(self, *args, **kwargs):
        res = _cali_zero(
            self.cxn, self.LO_Server,
            self.dv, self.reg, self.fpga, *args, **kwargs)
        return res

    def scan_spectrum(self, f_center, f_span):
        """get data from spectrum_analyzer
        Returns:
            (freqs, dBs)
            freqs: frequency points
            dBs: power points in unit of dB
        """
        (freqs, dBs)= calibrate.scan_spectrum(
            self.cxn, self.reg, f_center, f_span)
        return (freqs, dBs)
